"""
# 生成k-均值示例
"""

import numpy as np
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings('ignore')

n_data = 1000
seed = 1
n_clusters = 4

# 生成符合高斯随机分布的点团，并运行k-均值算法
blobs, blob_labels = make_blobs(n_samples=n_data, n_features=2,
                                centers=n_clusters, random_state=seed)
clusters_blob = KMeans(n_clusters=n_clusters, random_state=seed).fit_predict(blobs)

# 生成完全随机的数据，并运行k-均值算法
uniform = np.random.rand(n_data, 2)
clusters_uniform = KMeans(n_clusters=n_clusters, random_state=seed).fit_predict(uniform)

# 结果可视化
fig, ax = plt.subplots(2, 2)
axes = ax.flatten()

axes[0].scatter(blobs[:, 0], blobs[:, 1], c=blob_labels, cmap='gist_rainbow')
axes[0].set_title('(a) Four randomly generated blobs', fontsize=8)

axes[1].scatter(blobs[:, 0], blobs[:, 1], c=clusters_blob, cmap='gist_rainbow')
axes[1].set_title('(b) Clusters found via K-means', fontsize=8)

axes[2].scatter(uniform[:, 0], uniform[:, 1], cmap='gist_rainbow')
axes[2].set_title('(c) 1000 randomly generated points', fontsize=8)

axes[3].scatter(uniform[:, 0], uniform[:, 1], c=clusters_uniform, cmap='gist_rainbow')
axes[3].set_title('(d) Clusters found via K-means', fontsize=8)

fig.savefig('可视化/k-均值算法.png')
fig.show()
